package com.zf.experiment.limit;

import com.zf.experiment.annotation.AopAccessLimit;
import com.zf.experiment.constants.BasicConstants;
import com.zf.experiment.enums.AccessLimitMode;
import com.zf.experiment.exception.BusinessException;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.redisson.api.RLock;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.TimeUnit;

/**
 * 通过 AOP的方式 使用 AopAccessLimit 注解限流
 *
 * @author Zhou Feng
 */
@Slf4j
@Aspect
@Component
public class AopAccessLimitAspect {

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;
    @Autowired
    private RedissonClient redissonClient;

    //切入点：增强所有添加 该注解的类或者方法
    @Pointcut("@annotation(com.zf.experiment.annotation.AopAccessLimit)")
    public void testLimit() {
    }

    /**
     * 使用 redisson 作为分布式锁，进行限流
     * @param jp
     */
    // @Before("testLimit()")
    public void doLimit(JoinPoint jp) {
        Signature signature = jp.getSignature();
        Method method;
        if (!(signature instanceof MethodSignature)) {
            return;
        }
        method = ((MethodSignature) signature).getMethod();
        AopAccessLimit aopAccessLimit = method.getAnnotation(AopAccessLimit.class);
        if (aopAccessLimit == null || aopAccessLimit.seconds() <= 0 || aopAccessLimit.maxCount() < 0) {
            return;
        }
        int seconds = aopAccessLimit.seconds();
        int maxCount = aopAccessLimit.maxCount();
        String key = getLimitKey(aopAccessLimit, method);
        Object tmpNumber = redisTemplate.opsForValue().get(key);
        // 限流key不存在，则只能有一个线程来生成新的key并设置过期时间
        if (tmpNumber == null) {
            String lockKey = BasicConstants.REDISSON_LOCK_PREFIX + method.getDeclaringClass().getName() + "-" + method.getName();
            RLock lock = redissonClient.getLock(lockKey);
            lock.lock();
            try {
                // 再次判断限流key是否存在
                if (redisTemplate.opsForValue().get(key) != null){
                    //限流key存在，则直接递增 1
                    tmpNumber = redisTemplate.opsForValue().increment(key).intValue();
                }else {
                    // 不存在就设置限流key的值为 1
                    redisTemplate.opsForValue().set(key, BasicConstants.START_NUMBER, seconds, TimeUnit.SECONDS);
                    tmpNumber = BasicConstants.START_NUMBER;
                }
            } finally {
                lock.unlock();
            }
        } else {
            //限流key存在，则直接递增 1
            tmpNumber = redisTemplate.opsForValue().increment(key).intValue();
        }
        log.info("{}", tmpNumber);
        //判断是否需要限流
        if (maxCount < (int) tmpNumber) {
            throw new BusinessException("限流了.....");
        }
        log.info("没有限流");
    }


    /**
     * 通过lua脚本执行限流操作, 推荐使用lua
     */
    @Before("testLimit()")
    public void doLimitByLua(JoinPoint jp) {
        Signature signature = jp.getSignature();
        Method method;
        if (!(signature instanceof MethodSignature)) {
            return;
        }
        method = ((MethodSignature) signature).getMethod();
        AopAccessLimit aopAccessLimit = method.getAnnotation(AopAccessLimit.class);
        if (aopAccessLimit == null || aopAccessLimit.seconds() <= 0 || aopAccessLimit.maxCount() < 0) {
            return;
        }
        int seconds = aopAccessLimit.seconds();
        int maxCount = aopAccessLimit.maxCount();
        String key = getLimitKey(aopAccessLimit, method);

        //读取 lua 脚本
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        redisScript.setLocation(new ClassPathResource("lua/redis_limit.lua"));
        //必须设置
        redisScript.setResultType(Long.class);
        // redis中的key
        List<String> keys = Collections.singletonList(key);
        //执行 lua脚本，获取当前的访问数量
        Long tmpNumber = Optional.ofNullable(redisTemplate.execute(redisScript, keys, maxCount, seconds)).orElse(0L);
        log.info("{}", tmpNumber);
        //判断是否需要限流
        if (maxCount < tmpNumber.intValue()) {
            throw new BusinessException("限流了.....");
        }
        log.info("没有限流");
    }


    public String getLimitKey(AopAccessLimit accessLimit, Method method) {
        StringBuilder sb = new StringBuilder(accessLimit.prefix());
        // 按IP限流
        if (accessLimit.mode() == AccessLimitMode.IP) {
            ServletRequestAttributes requestAttributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
            String remoteAddr = requestAttributes.getRequest().getRemoteAddr();
            sb.append(remoteAddr).append(":");
        }
        // 按 访问数量限流
        Class<?> declaringClass = method.getDeclaringClass();
        sb.append(declaringClass.getName()).append("-").append(method.getName());
        return sb.toString();
    }
}
